package proxy

import (
	"context"
	"net/http"
	"strings"
	"sync"
	"time"

	"go-gateworker/internal/config"

	"golang.org/x/time/rate"
)

// PathLimiter 路径限流器
type PathLimiter struct {
	Path    string
	Limiter *rate.Limiter
}

// RateLimiter 限流器
type RateLimiter struct {
	globalLimiter *rate.Limiter
	pathLimiters  []*PathLimiter
	mu            sync.RWMutex
}

// NewRateLimiter 创建新的限流器
func NewRateLimiter(limits []config.RateLimit) *RateLimiter {
	rl := &RateLimiter{
		// 默认全局限流器，每秒1000个请求，突发容量100
		globalLimiter: rate.NewLimiter(1000, 100),
		pathLimiters:  make([]*PathLimiter, 0, len(limits)),
	}

	// 添加路径限流器
	for _, limit := range limits {
		rl.pathLimiters = append(rl.pathLimiters, &PathLimiter{
			Path:    limit.Path,
			Limiter: rate.NewLimiter(rate.Limit(limit.RequestsPerSec), limit.BurstSize),
		})
	}

	return rl
}

// UpdateGlobalLimit 更新全局限流配置
func (rl *RateLimiter) UpdateGlobalLimit(requestsPerSec, burstSize int) {
	rl.mu.Lock()
	defer rl.mu.Unlock()

	// 创建新的全局限流器
	rl.globalLimiter = rate.NewLimiter(rate.Limit(requestsPerSec), burstSize)
}

// UpdateLimits 更新限流配置
func (rl *RateLimiter) UpdateLimits(limits []config.RateLimit) {
	rl.mu.Lock()
	defer rl.mu.Unlock()

	// 清空现有路径限流器
	rl.pathLimiters = make([]*PathLimiter, 0, len(limits))

	// 添加路径限流器
	for _, limit := range limits {
		rl.pathLimiters = append(rl.pathLimiters, &PathLimiter{
			Path:    limit.Path,
			Limiter: rate.NewLimiter(rate.Limit(limit.RequestsPerSec), limit.BurstSize),
		})
	}
}

// Allow 检查请求是否允许通过
func (rl *RateLimiter) Allow(path string) bool {
	rl.mu.RLock()
	defer rl.mu.RUnlock()

	// 检查路径限流
	for _, pl := range rl.pathLimiters {
		if strings.HasPrefix(path, pl.Path) {
			if !pl.Limiter.Allow() {
				return false
			}
			break
		}
	}

	// 检查全局限流
	return rl.globalLimiter.Allow()
}

// Wait 等待直到可以通过限流或超时
func (rl *RateLimiter) Wait(path string, timeout time.Duration) bool {
	ctx, cancel := rl.getTimeoutContext(timeout)
	defer cancel()

	rl.mu.RLock()
	defer rl.mu.RUnlock()

	// 检查路径限流
	for _, pl := range rl.pathLimiters {
		if strings.HasPrefix(path, pl.Path) {
			if err := pl.Limiter.Wait(ctx); err != nil {
				return false
			}
			break
		}
	}

	// 检查全局限流
	return rl.globalLimiter.Wait(ctx) == nil
}

// getTimeoutContext 获取带超时的上下文
func (rl *RateLimiter) getTimeoutContext(timeout time.Duration) (context.Context, context.CancelFunc) {
	// 使用标准context包创建带超时的上下文
	return context.WithTimeout(context.Background(), timeout)
}

// RateLimitMiddleware 限流中间件
func RateLimitMiddleware(next http.Handler, limiter *RateLimiter) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if !limiter.Allow(r.URL.Path) {
			w.WriteHeader(http.StatusTooManyRequests)
			w.Write([]byte("Too many requests"))
			return
		}
		next.ServeHTTP(w, r)
	})
}
